﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using Microsoft.Cci;
using Microsoft.Cci.MetadataReader;
using Microsoft.Cci.MutableCodeModel;
using System.Runtime.Serialization; // needed for defining exception .ctors

namespace HeapProfiler
{
	public class Rewriter
	{
		private static String RuntimeDLL;
		private static String RuntimeDLLFile;
		private static String RuntimeDLL2;
		private static String RuntimeDLLFile2;
		
		static Rewriter()
		{
			String currentdir = System.IO.Path.GetDirectoryName(System.Reflection.Assembly.GetEntryAssembly().Location);
			RuntimeDLL = System.IO.Path.Combine(currentdir, @"RuntimeSupport.dll");
			RuntimeDLLFile = @"RuntimeSupport.dll";
			RuntimeDLL2 = System.IO.Path.Combine(currentdir, @"HeapModel.dll");
			RuntimeDLLFile2 = @"HeapModel.dll";
		}

		public static void DoRewrite(String exefile, String optionalrestrict)
		{
			var host = new PeReader.DefaultHost();

			IModule module = host.LoadUnitFrom(exefile) as IModule;
			if(module == null || module == Dummy.Module || module == Dummy.Assembly)
			{
				Console.WriteLine(exefile + " is not a PE file containing a CLR assembly, or an error occurred when loading it.");
				return;
			}

			String pdbFile = System.IO.Path.ChangeExtension(module.Location, "pdb");
			if(!System.IO.File.Exists(pdbFile))
			{
				Console.WriteLine("Could not load the PDB file for '" + module.Name.Value + "'.");
				return;
			}

			using(System.IO.Stream pdbStream = System.IO.File.OpenRead(pdbFile))
			{
				PdbReader pdbReader = new PdbReader(pdbStream, host);

				String copyFromRuntime = RuntimeDLL;
				String runtimeSupportLoc = System.IO.Path.Combine(System.IO.Path.GetDirectoryName(module.Location), RuntimeDLLFile);
				if(System.IO.File.Exists(runtimeSupportLoc))
					System.IO.File.Delete(runtimeSupportLoc);
				System.IO.File.Copy(copyFromRuntime, runtimeSupportLoc);

				String copyFromRuntime2 = RuntimeDLL2;
				String runtimeSupportLoc2 = System.IO.Path.Combine(System.IO.Path.GetDirectoryName(module.Location), RuntimeDLLFile2);
				if(System.IO.File.Exists(runtimeSupportLoc2))
					System.IO.File.Delete(runtimeSupportLoc2);
				System.IO.File.Copy(copyFromRuntime2, runtimeSupportLoc2);

				IModule supportMod = host.LoadUnitFrom(runtimeSupportLoc) as IModule;
				if(supportMod == null || supportMod == Dummy.Assembly)
				{
					Console.WriteLine("Could not load the RuntimeSupport file for '" + supportMod.Name.Value + "'.");
					return;
				}

				////////////////////
				//Rewrite File
				ILRewrite mutator = new ILRewrite(host, pdbReader, supportMod as IAssembly, optionalrestrict);
				module = mutator.Visit(module);

				String basepath = System.IO.Path.ChangeExtension(module.Location, null);
				String outputpath = basepath + "_prof.exe";
				String pdbpath = basepath + "_prof.pdb";

				if(System.IO.File.Exists(outputpath))
					System.IO.File.Delete(outputpath);

				if(System.IO.File.Exists(pdbpath))
					System.IO.File.Delete(pdbpath);

				// Need to not pass in a local scope provider until such time as we have one that will use the mutator
				// to remap things (like the type of a scope constant) from the original assembly to the mutated one.
				using(var pdbWriter = new PdbWriter(pdbpath, pdbReader))
					PeWriter.WritePeToStream(module, host, System.IO.File.Create(outputpath), pdbReader, pdbReader, pdbWriter);

				pdbReader = null;
			}

			host = null;
			module = null;

			return; // success
		}
	}
}
